from sentence_transformers import SentenceTransformer
from pinecone import Pinecone
from tqdm import tqdm

## Desc Model
model = SentenceTransformer("F:/BigModel/sentence-transformers-model/BAAI_bge-large-zh-v1.5", device="cuda")


def encode_sentence(sentence):
    sentence_embedding = model.encode(sentence, normalize_embeddings=True)

    return sentence_embedding


if __name__ == "__main__":
    # sentence = "你是谁"
    # print("输入的文本特征为：\n", sentence)
    # embedding = encode_sentence(sentence)
    # print("输出的向量特征为：\n", embedding)
    import pickle

    with open("LLM_modeled_res.pkl", "rb") as f:
        llm_res = pickle.load(f)

    llm_embeddings = {}
    for sid, llm_res in tqdm(llm_res.items()):
        llm_embeddings[sid] = encode_sentence(llm_res).tolist()

    with open("llm_embedding.pkl", "wb") as f:
        pickle.dump(llm_embeddings, f)